#!/usr/bin/env python3
import os
import sys
import argparse
import subprocess
import json
import pandas as pd
from pathlib import Path
from douyin_processor import DouyinProcessor


class DialectProcessor:
    def __init__(self, output_dir="dialect_audio"):
        """Initialize the dialect processor with output directory"""
        self.output_dir = output_dir
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            
        # Create subdirectory for each dialect streamer
        self.streamers_dir = os.path.join(output_dir, "streamers")
        if not os.path.exists(self.streamers_dir):
            os.makedirs(self.streamers_dir)
            
        # Initialize the DouyinProcessor for processing videos
        self.douyin_processor = DouyinProcessor(output_dir)
        
        # Dictionary of Chinese dialects and their regions
        self.dialects = {
            "北京话": "Beijing",
            "上海话": "Shanghai",
            "粤语": "Cantonese",
            "四川话": "Sichuan",
            "东北话": "Northeastern",
            "湖南话": "Hunan",
            "河南话": "Henan",
            "陕西话": "Shaanxi",
            "山东话": "Shandong",
            "闽南话": "Min Nan",
            "客家话": "Hakka",
            "湖北话": "Hubei",
            "安徽话": "Anhui", 
            "江西话": "Jiangxi",
            "广西话": "Guangxi",
            "云南话": "Yunnan",
            "贵州话": "Guizhou",
            "甘肃话": "Gansu",
            "河北话": "Hebei",
            "山西话": "Shanxi"
        }
    
    def process_streamer(self, streamer_info):
        """Process a single streamer's videos
        
        Args:
            streamer_info: Dictionary with keys:
                - name: Streamer's name
                - dialect: Dialect spoken
                - urls: List of video URLs
                
        Returns:
            List of segment files
        """
        name = streamer_info['name']
        dialect = streamer_info['dialect']
        urls = streamer_info['urls']
        
        print(f"Processing {name} ({dialect}) with {len(urls)} videos...")
        
        # Create directory for this streamer
        streamer_dir = os.path.join(self.streamers_dir, f"{name}_{dialect}")
        if not os.path.exists(streamer_dir):
            os.makedirs(streamer_dir)
            
        all_segments = []
        for i, url in enumerate(urls, 1):
            output_name = f"{name}_{dialect}_{i}"
            segments = self.douyin_processor.process_video(url, output_name)
            
            if segments:
                # Copy the first 5 segments (or all if less) to streamer directory
                n_segs = min(5, len(segments))
                for j in range(n_segs):
                    src = segments[j]
                    dst = os.path.join(streamer_dir, os.path.basename(src))
                    try:
                        import shutil
                        shutil.copy2(src, dst)
                        all_segments.append(dst)
                        print(f"Copied segment to {dst}")
                    except Exception as e:
                        print(f"Error copying segment: {e}")
        
        print(f"Finished processing {name}, saved {len(all_segments)} segments")
        return all_segments
    
    def load_streamers_from_file(self, file_path):
        """Load streamers information from JSON file"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                streamers = json.load(f)
            return streamers
        except Exception as e:
            print(f"Error loading streamers file: {e}")
            return []
    
    def save_streamers_to_file(self, streamers, file_path):
        """Save streamers information to JSON file"""
        try:
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(streamers, f, ensure_ascii=False, indent=2)
            print(f"Saved streamers info to {file_path}")
        except Exception as e:
            print(f"Error saving streamers file: {e}")
    
    def process_all_streamers(self, streamers):
        """Process all streamers in the list"""
        results = []
        for streamer in streamers:
            segments = self.process_streamer(streamer)
            if segments:
                results.append({
                    'name': streamer['name'],
                    'dialect': streamer['dialect'],
                    'segments': segments
                })
        
        # Save summary to CSV
        summary_file = os.path.join(self.output_dir, "streamers_summary.csv")
        rows = []
        for result in results:
            for segment in result['segments']:
                rows.append({
                    'Streamer': result['name'],
                    'Dialect': result['dialect'],
                    'Segment_File': segment
                })
        
        if rows:
            df = pd.DataFrame(rows)
            df.to_csv(summary_file, index=False, encoding='utf-8')
            print(f"Saved summary to {summary_file}")
        
        return results


def main():
    parser = argparse.ArgumentParser(description="Process dialect audio from streamers")
    parser.add_argument("--streamers_file", help="JSON file with streamers information")
    parser.add_argument("--output_dir", default="dialect_audio", help="Output directory")
    
    args = parser.parse_args()
    
    processor = DialectProcessor(args.output_dir)
    
    if args.streamers_file:
        streamers = processor.load_streamers_from_file(args.streamers_file)
        if streamers:
            processor.process_all_streamers(streamers)
        else:
            print("No streamers found in the file or file could not be loaded")
    else:
        # Example streamers data - this would typically be loaded from a file
        # This is just a placeholder and should be replaced with actual data
        print("No streamers file provided. Please create a JSON file with streamers information.")
        print("Example format:")
        example = [
            {
                "name": "streamer1",
                "dialect": "四川话",
                "urls": ["https://v.douyin.com/example1/", "https://v.douyin.com/example2/"]
            },
            {
                "name": "streamer2",
                "dialect": "粤语",
                "urls": ["https://v.douyin.com/example3/"]
            }
        ]
        print(json.dumps(example, ensure_ascii=False, indent=2))
        
        # Create example file
        example_file = os.path.join(args.output_dir, "example_streamers.json")
        processor.save_streamers_to_file(example, example_file)
        print(f"Created example streamers file at {example_file}")


if __name__ == "__main__":
    main() 